

#include "yolov5.h"
using namespace std;
using namespace cv;
using namespace cv::dnn;
#include "utils.h"
#include "./timer.h"
#include <unistd.h>
#include <stdio.h>

void datap::reset()
{
    classIds.resize(0);
    confidences.resize(0);
    boxes.resize(0);
    nms_result.resize(0);
}
void datap::press(float *pdata, size_t dlen, int net_width, float conf_thres, float nmsconf, float *sc)
{
    reset();
    // if (DEBUG)
    // {
    //     cout<<dlen<<" "<<net_width<<endl;
    // }
    // dlen = dlen / net_width;
    for (int i = 0; i < dlen; i++)
    {
        int index = i * net_width;
        if (pdata[index + 4] > conf_thres)
        {
            float *max_cls_pos = std::max_element(pdata + index + 5, pdata + index + net_width);
            // {

            float conf = pdata[index + 4] * (*max_cls_pos);
            float x = pdata[index + 0] * sc[0]; // x
            float y = pdata[index + 1] * sc[1]; // y
            float w = pdata[index + 2] * sc[0]; // w
            float h = pdata[index + 3] * sc[1]; // h
            int left = (x - 0.5 * w);
            int top = (y - 0.5 * h);
            confidences.push_back(conf);
            classIds.push_back(int(max_cls_pos - (pdata + index + 5)));

            boxes.push_back(Rect(left, top, int(w), int(h)));
            // }
        }
    }
    cv::dnn::NMSBoxes(boxes, confidences, conf_thres, nmsconf, nms_result);
}
void datap::print()
{
    std::cout << boxes.size() << " " << nms_result.size() << endl;
    for (auto &it : nms_result)
    {
        std::cout << boxes[it] << " " << confidences[it] << " " << classIds[it] << std::endl;
    }
}
void datap::getbox(vector<Output> &output)
{
    for (int i = 0; i < nms_result.size(); i++)
    {
        int idx = nms_result[i];
        Output result;
        result.id = classIds[idx];
        result.confidence = confidences[idx];
        result.box = boxes[idx];
        output.push_back(result);
    }
}
void datap::getbox(vector<Outputf> &output)
{
    for (int i = 0; i < nms_result.size(); i++)
    {
        int idx = nms_result[i];
        Outputf result;
        result.id = (float)classIds[idx];
        result.confidence = confidences[idx];
        result.box = boxes[idx];

        // cout << rsesult.box << endl;
        output.push_back(result);
    }
}
yolov5::~yolov5()
{
    delete dp;
}

yolov5::yolov5(string netPath, float _thresh) : thresh(_thresh)
{

    string fpath = cv::format("../config/config.yaml");
    ACC(fpath.c_str());
    cv::FileStorage fs((fpath), cv::FileStorage::READ);
    // system(string("cat " + fpath).c_str());
    // std::cout << CURRENT_VERSION << std::endl;
    string path;
    fs["mode_path"] >> path;
    fs["yolov5_P6"] >> yolov5_P6;
    fs["DEBUG"] >> DEBUG;
    fs["use_gpu"] >> cuda;
    fs["class_num"] >> class_num;
    className.resize(class_num);
    if (yolov5_P6)
    {
        netWidth = 1280;  // ONNX图片输入宽度
        netHeight = 1280; // ONNX图片输入高度
    }
    fs.release();
    std::cout << "yolov5_P6::::::" << yolov5_P6 << " " << netHeight << " " << netWidth << endl;
    try
    {
        cout << "-----load----" << path << std::endl;
        if (access(path.c_str(), F_OK) != 0)
            system("python3 /project/train/yolov5/exp.py");

        net = readNet(path);
    }
    catch (const std::exception &)
    {
        return;
    }
    // cuda
    if (cuda != 0)
    {
        net.setPreferableBackend(cv::dnn::DNN_BACKEND_CUDA);
        if (cuda == 1)
            net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA_FP16); // DNN_TARGET_CUDA DNN_TARGET_CUDA_FP16
        else
            net.setPreferableTarget(cv::dnn::DNN_TARGET_CUDA); // DNN_TARGET_CUDA DNN_TARGET_CUDA_FP16
    }
    // cpu
    else
    {
        net.setPreferableBackend(cv::dnn::DNN_BACKEND_DEFAULT);
        net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
    }
    dp = new datap();
    return;
}
void matp(cv::Mat mat)
{
    printf("mat.step[0]=%ld\n", mat.step[0]);                                                //每一行所占字节数目
    printf("mat.step[1]=%ld\n", mat.step[1]);                                                //每个元素所占字节数
    cout << "L::" << __LINE__ << " size:" << mat.size() << " " << mat.step[0] / (4) << endl; //每个元素所占字节数
}
// new_shape: h,w
cv::Mat letterbox(cv::Mat im, vector<int> new_shape, cv::Vec3b color, bool auto_ = true, bool scaleFill = false, bool scaleup = true, int stride = 32)
{
    // Resize and pad image while meeting stride - multiple constraints
    //  # current shape [height, width]
    //     if isinstance(new_shape, int):
    //         new_shape = (new_shape, new_shape)
    vector<float> shape{(float)im.rows, (float)im.cols};
    // #Scale ratio(new / old)
    float r = std::min(new_shape[0] / shape[0], new_shape[1] / shape[1]);

    if (!scaleup) //:  # only scale down, do not scale up (for better val mAP)
        r = std::min(r, 1.0f);
    // #Compute padding
    vector<float> ratio = {r, r};                                           //# width, height ratios
    int new_unpad[] = {int(round(shape[1] * r)), int(round(shape[0] * r))}; // w  h
    int dw = new_shape[1] - new_unpad[0], dh = new_shape[0] - new_unpad[1]; //#wh padding
    if (auto_)                                                              //:  # minimum rectangle
        dw = (dw % stride), dh = (dh % stride);                             //# wh padding
    else if (scaleFill)                                                     //:  # stretch
    {
        dw = dh = 0;
        new_unpad[0] = new_shape[1], new_unpad[1] = new_shape[0];
        ratio[0] = new_shape[1] / shape[1], ratio[1] = new_shape[0] / shape[0]; // #width, height ratios;
    }
    dw /= 2; // #divide padding into 2 sides
    dh /= 2;
    cv::Size ns{new_unpad[0], new_unpad[1]};
    if (int(shape[0]) != new_unpad[0] || (int(shape[1]) != new_unpad[1])) //  :  # resize
        cv::resize(im, im, ns, cv::INTER_LINEAR);
    int top = int(round(dh - 0.1)), bottom = int(round(dh + 0.1));
    int left = int(round(dw - 0.1)), right = int(round(dw + 0.1));
    cv::copyMakeBorder(im, im, top, bottom, left, right, cv::BORDER_CONSTANT, color); // #add border
    return im;
}

bool yolov5::operator()(Mat &SrcImg, vector<Output> &output, float _thresh)
{
    Mat blob;
    int col = SrcImg.cols;
    int row = SrcImg.rows;
    int maxLen = MAX(col, row);
    float sc[] = {(float)row / netHeight, (float)col / netWidth};
    auto netInputImg = letterbox(SrcImg, {1280, 1280}, cv::Vec3b(114, 114, 114));
    if (DEBUG)
    {
        matp(blob);
        matp(netInputImg);
    }
    blobFromImage(netInputImg, blob, 1 / 255.0, cv::Size(netWidth, netHeight), cv::Scalar(0, 0, 0), true, false);

    //如果在其他设置没有问题的情况下但是结果偏差很大，可以尝试下用下面两句语句
    // blobFromImage(netInputImg, blob, 1 / 255.0, cv::Size(netWidth, netHeight), cv::Scalar(104, 117, 123), true, false);
    // blobFromImage(netInputImg, blob, 1 / 255.0, cv::Size(netWidth, netHeight), cv::Scalar(114, 114,114), true, false);
    net.setInput(blob);
    std::vector<cv::Mat> netOutputImg;
    // {
    //     if (DEBUG)
    //         Timer t("a");
    net.forward(netOutputImg, net.getUnconnectedOutLayersNames());
    // }
    std::vector<int> classIds;      //结果id数组
    std::vector<float> confidences; //结果每个id对应置信度数组
    std::vector<cv::Rect> boxes;    //每个id矩形框
    float ratio_h = (float)netInputImg.rows / netHeight;
    float ratio_w = (float)netInputImg.cols / netWidth;
    // className.resize(80);
    // className.resize(1);
    float *pdata = (float *)netOutputImg[0].data;
    cv::Size outsize = netOutputImg[0].size();
    int net_width = netOutputImg[0].step[1] / sizeof(float); // + 5; //输出的网络宽度是类别数+5

    dp->press(pdata, outsize.width, net_width, _thresh, nmsThreshold, sc);
    dp->getbox(output);

    if (DEBUG)
    {
        cout << netOutputImg[0].size() << " " << netOutputImg.size() << " " << netOutputImg[0].channels() << netOutputImg[0].depth() << endl;
        matp(netOutputImg[0]);
        cout << "WWWWWWWWWW=" << net_width << endl;
        dp->print();
    }

    if (output.size())
        return true;
    else
        return false;
}

void yolov5::drawPred(Mat &img, vector<Output> result, vector<Scalar> color)
{
    for (int i = 0; i < result.size(); i++)
    {
        int left, top;
        left = result[i].box.x;
        top = result[i].box.y;
        int color_num = i;
        rectangle(img, result[i].box, color[result[i].id], 2, 8);

        string label = className[result[i].id] + ":" + to_string(result[i].confidence);

        int baseLine;
        Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
        top = max(top, labelSize.height);
        // rectangle(frame, Point(left, top - int(1.5 * labelSize.height)), Point(left + int(1.5 * labelSize.width), top + baseLine), Scalar(0, 255, 0), FILLED);
        putText(img, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 1, color[result[i].id], 2);
    }
    // imshow("1", img);
    // imwrite("out.bmp", img);
    // waitKey();
    // destroyAllWindows();
}
